from enum import Enum
import os
import numpy as np
from torch.utils.data import Dataset

class TDMaterial(Enum):
    Graphene = 'Graphene'
    MoS2 = 'MoS2'
    WTe2 = 'WTe2'
    BN = 'BN'
    Ensemble = 'Ensemble_G_BN_WTe2_MoS2'
    
    def __str__(self):
        return self.value

class Split(Enum):
    TRAIN = 'train'
    VAL = 'val'
    
    def __str__(self):
        return self.value

class TDMaterials_dataset(Dataset):
    def __init__(self, base_dir, split=Split.TRAIN, transform=None, tdmaterial=TDMaterial.Graphene):
        self.transform = transform  # using transform in torch!
        self.split = split
        self.data_dir = os.path.join(base_dir,f'{tdmaterial}_{split}.npz')
            
        self.prepare()

    def __len__(self):
        return len(self.image_info)

    def __getitem__(self, idx):
        image = self.image[idx]
        label = self.label[idx]
        if self.transform:
            image, label = self.transform((image, label))
        return image, label
    
        
    def prepare(self):
        """Prepares the Dataset class for use.
        """
        # load npz file
        data = np.load(self.data_dir, allow_pickle=True)
        self.image = data['image']
        self.label = data['label']
        self.image_name = data['image_name']
        self.class_names = data['class_names']
        self.num_classes = data['num_classes']
        self.image_info = data['image_info']
        self.class_info = data['class_info']
        
        self.class_ids = np.arange(self.num_classes)
        self.num_images = len(self.image_info)
                            